import torch
import torchvision
import torchvision.transforms as transforms
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, classification_report
import matplotlib.pyplot as plt
import numpy as np
import os

# 加载FashionMNIST完整数据集
transform = transforms.Compose([transforms.ToTensor()])

train_dataset = torchvision.datasets.FashionMNIST(
    root='/workspace/FashionMNIST', train=True, transform=transform, download=True)
test_dataset = torchvision.datasets.FashionMNIST(
    root='/workspace/FashionMNIST', train=False, transform=transform, download=True)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=len(train_dataset), shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=len(test_dataset), shuffle=False)

# 一次性加载所有训练数据
print("加载完整训练数据...")
for images, labels in train_loader:
    X_train = images.view(images.shape[0], -1).numpy()
    y_train = labels.numpy()

# 一次性加载所有测试数据
print("加载完整测试数据...")
for images, labels in test_loader:
    X_test = images.view(images.shape[0], -1).numpy()
    y_test = labels.numpy()

# 定义SVC分类器 (Poly kernel, C=10)
svc = SVC(kernel='poly', C=10, verbose=1)

print("开始训练SVC（完整数据集）...")
svc.fit(X_train, y_train)

# 进行预测
print("开始预测...")
y_pred = svc.predict(X_test)

# 输出性能评估
accuracy = accuracy_score(y_test, y_pred)
print("完整测试集准确率:", accuracy)
print("详细分类报告:")
print(classification_report(y_test, y_pred))

def visualize_predictions(X_test, y_test, y_pred, num_samples=10):
    """
    可视化测试集中的样本及其预测结果，并保存图像
    
    参数:
    - X_test: 测试数据特征
    - y_test: 测试数据真实标签
    - y_pred: 模型预测的标签
    - num_samples: 要可视化的样本数量
    """
    # 设置保存路径为FashionMNIST下的output文件夹
    save_dir = '/workspace/FashionMNIST/output'
    
    # FashionMNIST的类别名称
    class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
                   'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
    
    # 随机选择num_samples个样本
    indices = np.random.choice(len(y_test), num_samples, replace=False)
    
    # 创建一个网格图
    fig = plt.figure(figsize=(12, 8))
    
    # 使用英文标题
    fig.suptitle("SVC Model Predictions on FashionMNIST Test Set", fontsize=16)
    
    for i, idx in enumerate(indices):
        # 创建子图
        plt.subplot(2, num_samples//2, i+1)
        
        # 将一维特征重塑为28x28的图像
        img = X_test[idx].reshape(28, 28)
        plt.imshow(img, cmap='gray')
        
        # 设置标题为预测结果，颜色根据预测是否正确而定
        title_color = 'green' if y_pred[idx] == y_test[idx] else 'red'
        
        # 使用英文标题
        title = f'Pred: {class_names[y_pred[idx]]}\nTrue: {class_names[y_test[idx]]}'
        plt.title(title, color=title_color)
        
        # 隐藏坐标轴
        plt.axis('off')
    
    plt.tight_layout()
    plt.subplots_adjust(top=0.9)  # 为总标题留出空间
    
    # 创建保存目录（如果不存在）
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    
    # 使用固定文件名保存图像（会覆盖同名文件）
    file_path = os.path.join(save_dir, "svc_fashion_mnist_predictions.png")
    
    # 保存图像
    plt.savefig(file_path, dpi=300, bbox_inches='tight')
    print(f"图像已保存到: {file_path}")
    
    # 显示图像
    plt.show()

# 调用绘图函数显示预测结果并保存图像
visualize_predictions(X_test, y_test, y_pred, num_samples=10)

from sklearn.decomposition import PCA  # Ensure this is imported at the top of your file

def visualize_pca_results(X_train, X_test, y_train, y_test, y_pred):
    """
    使用PCA将高维特征降维到2D并可视化SVC分类结果
    
    Parameters:
    - X_train: Training data features (numpy array)
    - X_test: Test data features (numpy array)
    - y_train: Training data labels (numpy array)
    - y_test: Test data true labels (numpy array)
    - y_pred: Model predicted labels (numpy array)
    """
    print("Starting PCA dimensionality reduction analysis...")

    # Initialize PCA with 3 components (we’ll use the first 2 for visualization)
    pca = PCA(n_components=3)
    
    # Fit PCA model using training data
    print("Fitting PCA model with training data...")
    pca.fit(X_train)
    
    # Record and display the explained variance ratio of the principal components
    explained_variance = pca.explained_variance_ratio_
    print(f"Variance explained by first two components: {explained_variance[0]:.4f}, {explained_variance[1]:.4f}")
    print(f"Cumulative explained variance: {sum(explained_variance):.4f}")
    
    # Transform test data to PCA space
    print("Transforming test data with PCA...")
    X_test_pca = pca.transform(X_test)
    
    # Set save directory
    save_dir = '/workspace/FashionMNIST/output'
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        print(f"Created output directory: {save_dir}")
    
    # Start drawing the PCA scatter plot
    print("Drawing PCA dimensionality reduction scatter plot...")
    
    # Define FashionMNIST class names
    class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
                   'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
    
    # Randomly sample a subset of test points for clarity
    np.random.seed(42)  # Set seed for reproducibility
    sample_indices = np.random.choice(len(y_test), min(2000, len(y_test)), replace=False)
    
    # Create figure
    plt.figure(figsize=(14, 10))
    
    # Define markers and colors for each class
    markers = ['o', 'v', '^', '<', '>', 's', 'p', 'h', 'D', '*']
    colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', 
              '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
    
    # Plot points for each class, separating correct and incorrect predictions
    for i, class_name in enumerate(class_names):
        # Get indices of samples belonging to this class from the sampled subset
        class_indices = np.where(y_test[sample_indices] == i)[0]
        if len(class_indices) > 0:
            actual_indices = sample_indices[class_indices]
            # Separate correct and incorrect predictions
            correct_mask = y_pred[actual_indices] == i
            correct_points = actual_indices[correct_mask]
            incorrect_points = actual_indices[~correct_mask]
            
            # Plot correct predictions
            if len(correct_points) > 0:
                plt.scatter(X_test_pca[correct_points, 0], X_test_pca[correct_points, 1], 
                            marker=markers[i], color=colors[i], s=50, alpha=0.6, 
                            label=f'{class_name}', edgecolors='black', linewidth=0.5)
            
            # Plot incorrect predictions as red 'x'
            if len(incorrect_points) > 0:
                plt.scatter(X_test_pca[incorrect_points, 0], X_test_pca[incorrect_points, 1], 
                            marker='x', color='red', s=40, alpha=0.8)
    
    # Add class centroids and labels
    for i, class_name in enumerate(class_names):
        class_indices = np.where(y_test == i)[0]
        if len(class_indices) > 0:
            center_x = np.mean(X_test_pca[class_indices, 0])
            center_y = np.mean(X_test_pca[class_indices, 1])
            # Plot centroid with larger marker
            plt.scatter(center_x, center_y, marker=markers[i], s=200, 
                        color=colors[i], edgecolor='black', linewidth=2.0)
            # Annotate with class name
            plt.annotate(class_name, (center_x, center_y), 
                         xytext=(8, 8), textcoords='offset points', 
                         fontsize=12, fontweight='bold', 
                         bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.7))
    
    # Add title and axis labels with explained variance
    plt.title('FashionMNIST Dataset PCA Visualization (SVC Classification)', fontsize=16)
    plt.xlabel(f'First Principal Component (Explained Variance: {explained_variance[0]:.2%})', fontsize=12)
    plt.ylabel(f'Second Principal Component (Explained Variance: {explained_variance[1]:.2%})', fontsize=12)
    
    # Add model accuracy annotation
    accuracy = accuracy_score(y_test, y_pred)
    plt.annotate(f'Model Accuracy: {accuracy:.4f}', 
                 xy=(0.02, 0.02), xycoords='axes fraction', fontsize=12, 
                 bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))
    
    # Add legend entry for incorrect predictions
    plt.scatter([], [], marker='x', color='red', s=40, label='Incorrect Predictions')
    
    # Add legend with multiple columns
    plt.legend(loc='upper center', fontsize=10, bbox_to_anchor=(0.5, -0.05), 
               ncol=5, frameon=True, fancybox=True, shadow=True)
    
    # Add grid and adjust layout
    plt.grid(alpha=0.3)
    plt.tight_layout()
    
    # Save the plot
    file_path = os.path.join(save_dir, "svc_fashion_mnist_pca_visualization.png")
    plt.savefig(file_path, dpi=300, bbox_inches='tight')
    print(f"PCA visualization image saved to: {file_path}")
    
    # Display the plot
    plt.show()
    
    # Print additional analysis: per-class accuracy and PCA dispersion
    print("PCA distribution by class:")
    for i, class_name in enumerate(class_names):
        class_indices = np.where(y_test == i)[0]
        if len(class_indices) > 0:
            correct_class = np.sum(y_pred[class_indices] == i)
            class_accuracy = correct_class / len(class_indices)
            class_std = np.std(X_test_pca[class_indices], axis=0)
            class_dispersion = np.mean(class_std)
            print(f"Class {class_name}: Accuracy={class_accuracy:.4f}, PCA Dispersion={class_dispersion:.4f}")

    print("PCA analysis completed.")

# Call the PCA visualization function
visualize_pca_results(X_train, X_test, y_train, y_test, y_pred)